package main

import (
	"context"
	"golang.org/x/time/rate"
	"io"
	"log"
	"net"
	"time"
)

// 常量配置
const (
	// 服务器监听端口
	LISTEN_PORT = ":9000"

	// 目标限速值 (Mbps)
	TARGET_RATE_MBPS = 6

	// 每秒字节限速值 (5Mbps = 5*1024*1024/8 bytes/s)
	BYTE_RATE_LIMIT = TARGET_RATE_MBPS * 1024 * 1024 / 8

	// TCP接收缓冲区大小 (8KB)
	// 通过控制TCP窗口大小实现内核层限速
	SO_RCVBUF_SIZE = 1 * 1024

	// 每次从连接读取的数据块大小 (4KB)
	READ_CHUNK_SIZE = 4 * 1024
)

// 程序入口点
func main() {
	// 输出服务器启动信息和配置参数
	log.Printf("Starting TCP rate limit server (target: %d Mbps)", TARGET_RATE_MBPS)
	log.Printf("Kernel buffer: %d KB | App chunk: %d KB", SO_RCVBUF_SIZE/1024, READ_CHUNK_SIZE/1024)

	// 创建TCP监听器
	listener, err := net.Listen("tcp", LISTEN_PORT)
	if err != nil {
		log.Fatalf("Listen failed: %v", err)
	}
	defer listener.Close()

	// 循环接受客户端连接
	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Printf("Accept error: %v", err)
			continue
		}

		// 为每个连接启动一个goroutine处理
		go handleConnection(conn)
	}
}

// 处理客户端连接
func handleConnection(conn net.Conn) {
	defer conn.Close()

	// 获取客户端地址信息
	clientAddr := conn.RemoteAddr().String()
	log.Printf("Client connected: %s", clientAddr)

	// ========== 内核层限速 ==========
	// 尝试设置TCP接收缓冲区大小，实现底层限速
	if tcpConn, ok := conn.(*net.TCPConn); ok {
		if err := setSocketRecvBuffer(tcpConn, SO_RCVBUF_SIZE); err != nil {
			log.Printf("WARN: kernel limit failed (%v), using app layer only", err)
		}
	}

	// ========== 应用层限速 ==========
	// 使用令牌桶算法进行应用层限速
	// rate.Limit(BYTE_RATE_LIMIT): 每秒产生BYTE_RATE_LIMIT个令牌
	// BYTE_RATE_LIMIT: 令牌桶容量，控制突发流量大小
	limiter := rate.NewLimiter(rate.Limit(BYTE_RATE_LIMIT), BYTE_RATE_LIMIT)

	// 创建读取缓冲区
	buf := make([]byte, READ_CHUNK_SIZE)

	// 添加速率统计相关变量
	var totalBytes int64       // 总接收字节数
	startTime := time.Now()    // 连接开始时间
	lastPrintTime := startTime // 上次打印速率的时间

	// 循环读取客户端数据
	for {
		// 等待足够的令牌才能继续读取
		// 令牌数量与要读取的数据块大小一致
		if err := limiter.WaitN(context.Background(), len(buf)); err != nil {
			log.Printf("Rate limit error: %v", err)
			return
		}

		// 从连接中读取数据
		n, err := conn.Read(buf)
		if err != nil {
			if err != io.EOF {
				log.Printf("Read error: %v", err)
			}
			break
		}

		// 累计接收的字节数
		totalBytes += int64(n)

		// 每秒打印一次当前速率
		now := time.Now()
		if now.Sub(lastPrintTime) >= time.Second {
			elapsed := now.Sub(startTime).Seconds()
			bytesPerSec := float64(totalBytes) / elapsed
			mbps := bytesPerSec * 8 / 1024 / 1024

			log.Printf("Client %s: Rate %.2f Mbps (%.2f KB/s)",
				clientAddr, mbps, bytesPerSec/1024)

			lastPrintTime = now
		}
	}

	// 连接关闭时打印总体统计信息
	elapsed := time.Since(startTime).Seconds()
	avgBytesPerSec := float64(totalBytes) / elapsed
	avgMbps := avgBytesPerSec * 8 / 1024 / 1024

	log.Printf("Client disconnected: %s | Avg rate: %.2f Mbps | Total: %.2f MB",
		clientAddr, avgMbps, float64(totalBytes)/(1024*1024))
}

// ========== 跨平台内核限速实现 ==========
// setSocketRecvBuffer函数已移至平台特定文件中:
// - socket_unix.go (Linux/macOS) - //go:build !windows
// - socket_windows.go (Windows) - //go:build windows
//
// 注意：不要在main.go中实现此函数，否则会导致重定义错误
// GoLand IDE可能无法识别构建标签，如遇到编译错误，请使用命令行构建：
// Windows: go build -o server.exe main.go socket_windows.go
// Linux/macOS: go build -o server main.go socket_unix.go
// 或使用交叉编译:
// CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -o tcpserver-linux main.go

// 仅在IDE中使用的默认实现，避免编译错误
func setSocketRecvBuffer(conn *net.TCPConn, size int) error {
	// 这是一个默认实现，当平台特定的实现不可用时使用
	// 实际的实现位于socket_windows.go和socket_unix.go中
	return nil
}
